from paddle.fluid.dygraph import Conv2D, BatchNorm, Dropout, Linear, Pool2D


def _addindent(s_, numSpaces):
    s = s_.split('\n')
    # don't do anything for single-line stuff
    if len(s) == 1:
        return s_
    first = s.pop(0)
    s = [(numSpaces * ' ') + line for line in s]
    s = '\n'.join(s)
    s = first + '\n' + s
    return s


def conv_repr(mod):
    s = ('{_num_channels}, {_num_filters}, kernel_size={kernel_size}'
         ', stride={_stride}')
    kernel_size = mod._filter_size
    if isinstance(kernel_size, int):
        kernel_size = [kernel_size, kernel_size]
    kernel_size = list(kernel_size)
    if mod._padding != [0] * len(mod._padding):
        s += ', padding={_padding}'
    if mod._dilation != [1] * len(mod._dilation):
        s += ', dilation={_dilation}'
    # if self.output_padding != [0] * len(self.output_padding):
    #     s += ', output_padding={output_padding}'
    if mod._groups is not None and mod._groups != 1:
        s += ', groups={_groups}'
    if mod._bias_attr is False:
        s += ', bias=False'
    return s.format(**mod.__dict__, kernel_size=kernel_size)


def bn_repr(mod):
    d = {
        "num_features": mod._parameters['weight'].shape[0],
        "eps": mod._epsilon,
        "momentum": mod._momentum,
        "affine": not mod._parameters['weight'].stop_gradient,
        "track_running_stats": True,
    }
    return '{num_features}, eps={eps}, momentum={momentum}, affine={affine}, ' \
           'track_running_stats={track_running_stats}'.format(**d)


def dropout_repr(mod):
    d = {
        "p": mod._dropout_prob
    }
    return 'p={p}'.format(**d)


def linear_repr(mod):
    s = 'in_features={in_features}, out_features={out_features}'
    d = {
        "in_features": mod._parameters['weight'].shape[0],
        "out_features": mod._parameters['weight'].shape[1],
    }
    if 'bias' not in mod._parameters:
        s += ', bias=False'
    return s.format(**d)


def pool2d_repr(mod):
    s = 'kernel_size={_pool_size}, stride={_pool_stride}'
    if mod._pool_padding != [0] * len(mod._pool_padding):
        s += ', padding={_pool_padding}'
    s += ", ceil_mode={_ceil_mode}, type={_pool_type}"
    return s.format(**mod.__dict__)


reg = {
    Conv2D: conv_repr,
    BatchNorm: bn_repr,
    Dropout: dropout_repr,
    Linear: linear_repr,
    Pool2D: pool2d_repr,
}


def layer_repr(mod):
    # We treat the extra repr like the sub-module, one item per line
    extra_lines = []

    if hasattr(mod, "extra_repr"):
        extra_repr = mod.extra_repr()
    else:
        extra_repr = reg.get(mod.__class__, lambda _: "")(mod)
    # empty string will be split into list ['']
    if extra_repr:
        extra_lines = extra_repr.split('\n')
    child_lines = []
    for key, module in mod._sub_layers.items():
        mod_str = layer_repr(module)
        mod_str = _addindent(mod_str, 2)
        child_lines.append('(' + key + '): ' + mod_str)
    lines = extra_lines + child_lines

    main_str = mod.__class__.__name__ + '('
    if lines:
        # simple one-liner info, which most builtin Modules will use
        if len(extra_lines) == 1 and not child_lines:
            main_str += extra_lines[0]
        else:
            main_str += '\n  ' + '\n  '.join(lines) + '\n'

    main_str += ')'
    return main_str